package com.xlh.util;

import com.jcraft.jsch.*;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;

import java.io.Closeable;
import java.io.IOException;
import java.io.InputStream;
import java.util.Vector;
import java.util.concurrent.TimeUnit;

/**
 * SSH工具
 *
 * @author cheer
 */
@Slf4j
public class SshUtil implements Closeable {

    private static final int SESSION_TIMEOUT = 30000;
    private Session session = null;

    private SshUtil(SSHInfo sshInfo) {
        try {
            JSch jsch = new JSch();
            session = jsch.getSession(sshInfo.getUser(), sshInfo.getHost(), sshInfo.getPort());
            session.setPassword(sshInfo.getPassword());
            session.setUserInfo(new MyUserInfo());
            session.connect(SESSION_TIMEOUT);
        } catch (JSchException e) {
            log.warn("host {} ssh failed: {}", sshInfo.getHost(), e.getMessage());
        }
    }

    public static SshUtil newInstance(String user, String password, String host, Integer port) {
        SSHInfo sshInfo = new SSHInfo(user, password, host, port);
        return new SshUtil(sshInfo);
    }

    public String exec(String cmd) throws IOException, JSchException, InterruptedException {
        ChannelExec channelExec = (ChannelExec) session.openChannel("exec");
        channelExec.setCommand(cmd);
        channelExec.setInputStream(null);
        channelExec.setErrStream(System.err);
        InputStream in = channelExec.getInputStream();
        channelExec.connect();

        StringBuilder buf = getStringBuilder(channelExec, in);
        channelExec.disconnect();
        return buf.toString();
    }

    public String execWithError(String cmd) throws IOException, JSchException, InterruptedException {
        ChannelExec channelExec = (ChannelExec) session.openChannel("exec");
        channelExec.setCommand(cmd);
        channelExec.setInputStream(null);
        channelExec.setErrStream(System.err);
        InputStream in = channelExec.getInputStream();
        channelExec.connect();

        InputStream err =  channelExec.getErrStream();
        StringBuilder errBuf = getStringBuilder(channelExec, err);
        if(StringUtils.isNotBlank(errBuf.toString())){
            channelExec.disconnect();
            return errBuf.toString();
        }

        StringBuilder buf = getStringBuilder(channelExec, in);
        channelExec.disconnect();
        return buf.toString();
    }

    private StringBuilder getStringBuilder(ChannelExec channelExec, InputStream in) throws IOException, InterruptedException {
        StringBuilder buf = new StringBuilder(1024);
        byte[] tmp = new byte[1024];
        while (true) {
            while (in.available() > 0) {
                int i = in.read(tmp, 0, 1024);
                if (i < 0) break;
                buf.append(new String(tmp, 0, i));
            }
            if (channelExec.isClosed()) {
                break;
            }
            TimeUnit.MILLISECONDS.sleep(100);
        }
        return buf;
    }

    public void execWithoutReturn(String cmd) throws JSchException {
        ChannelExec channelExec = (ChannelExec) session.openChannel("exec");
        channelExec.setCommand(cmd);
        channelExec.setInputStream(null);
        channelExec.setErrStream(System.err);
        channelExec.connect();
        channelExec.disconnect();
    }

    public void upload(String localFile, String remoteFile) throws JSchException, SftpException {
        ChannelSftp channelSftp = (ChannelSftp) session.openChannel("sftp");
        channelSftp.connect();

        channelSftp.put(localFile, remoteFile, ChannelSftp.OVERWRITE);
        channelSftp.disconnect();
    }

    public void download(String localFile, String remoteFile) throws JSchException, SftpException {
        ChannelSftp channelSftp = (ChannelSftp) session.openChannel("sftp");
        channelSftp.connect();

        channelSftp.get(localFile, remoteFile);
        channelSftp.disconnect();
    }

    public void get(String remoteFile, String localFile) throws JSchException, SftpException {
        ChannelSftp channelSftp = (ChannelSftp) session.openChannel("sftp");
        channelSftp.connect();

        channelSftp.get(remoteFile, localFile);
        channelSftp.disconnect();
    }

    public Vector<ChannelSftp.LsEntry> ls(String remoteFile) throws JSchException, SftpException {
        ChannelSftp channelSftp = (ChannelSftp) session.openChannel("sftp");
        channelSftp.connect();

        Vector<ChannelSftp.LsEntry> ls = channelSftp.ls(remoteFile);
        channelSftp.disconnect();
        return ls;
    }

    public Session getSession() {
        return session;
    }

    @Override
    public void close() {
        getSession().disconnect();
    }

    /**
     * SSH连接信息
     */
    public static class SSHInfo {
        private String user;
        private String password;
        private String host;
        private int port;

        public SSHInfo(String user, String password, String host, int port) {
            this.user = user;
            this.password = password;
            this.host = host;
            this.port = port;
        }

        public String getUser() {
            return user;
        }

        public String getPassword() {
            return password;
        }

        public String getHost() {
            return host;
        }

        public int getPort() {
            return port;
        }
    }

    /**
     * 自定义UserInfo
     */
    private static class MyUserInfo implements UserInfo {

        @Override
        public String getPassphrase() {
            return null;
        }

        @Override
        public String getPassword() {
            return null;
        }

        @Override
        public boolean promptPassword(String s) {
            return false;
        }

        @Override
        public boolean promptPassphrase(String s) {
            return false;
        }

        @Override
        public boolean promptYesNo(String s) {
            return true;
        }

        @Override
        public void showMessage(String s) {
            // 暂不实现
        }
    }
}
